import numpy as np
import numpy.fft as npf
import scipy.signal as sps
import misc

class UserError(Exception):
    def __init__(self, value):
        self.value = value
    def __str__(self):
        return repr(self.value)

class Baseband:
    def __init__(self,RB,fftlen):
        self.RB = RB
        self.fftlen = fftlen

        self.scs   = 15E3
        self.cp0   = 160./(self.scs*2048)
        self.cp16  = 144./(self.scs*2048)
        self.cps0  = int(self.cp0*self.fftlen*self.scs)
        self.cps16 = int(self.cp16*self.fftlen*self.scs)

    def modulator(self,x,mod='QPSK',skip=False):
        """ x is a list of binary integers (0,1). RB is the number of
        resource blocks and mod determines whether to user QPSK or 16QAM
        modulation.
        """
        
        "Check input"
        if mod != 'QPSK' and mod != '16QAM':
            raise UserError('Modulator must be QPSK or 16QAM')
        elif mod == 'QPSK' and len(x) % 24 != 0:
            raise UserError('Input must be dividable by 24*RB since QPSK is used.')
        elif mod == '16QAM' and (len(x) % 48 != 0) and skip == False:
            raise UserError('Input must be dividable by 48*RB since 16QAM is used.')

        "Perform modulation"
        if mod == 'QPSK':
            x_mod = []
            for i in range(len(x)/2):
                j = i*2
                if x[j]==0 and x[j+1]==0:
                    x_mod += [1./np.sqrt(2)+1j*1./np.sqrt(2)]
                elif x[j]==0 and x[j+1]==1:
                    x_mod += [1./np.sqrt(2)-1j*1./np.sqrt(2)]
                elif x[j]==1 and x[j+1]==0:
                    x_mod += [-1./np.sqrt(2)+1j*1./np.sqrt(2)]
                elif x[j]==1 and x[j+1]==1:
                    x_mod += [-1./np.sqrt(2)-1j*1./np.sqrt(2)]
        elif mod == '16QAM':
            x_mod = []
            for i in range(len(x)/4):
                j = i*4
                if x[j]==0 and x[j+1]==0 and x[j+2]==0 and x[j+3]==0:
                    x_mod += [1./np.sqrt(10)+1j*1./np.sqrt(10)]
                elif x[j]==0 and x[j+1]==0 and x[j+2]==0 and x[j+3]==1:
                    x_mod += [1./np.sqrt(10)+1j*3./np.sqrt(10)]
                elif x[j]==0 and x[j+1]==0 and x[j+2]==1 and x[j+3]==0:
                    x_mod += [3./np.sqrt(10)+1j*1./np.sqrt(10)]
                elif x[j]==0 and x[j+1]==0 and x[j+2]==1 and x[j+3]==1:
                    x_mod += [3./np.sqrt(10)+1j*3./np.sqrt(10)]
                elif x[j]==0 and x[j+1]==1 and x[j+2]==0 and x[j+3]==0:
                    x_mod += [1./np.sqrt(10)-1j*1./np.sqrt(10)]
                elif x[j]==0 and x[j+1]==1 and x[j+2]==0 and x[j+3]==1:
                    x_mod += [1./np.sqrt(10)-1j*3./np.sqrt(10)]
                elif x[j]==0 and x[j+1]==1 and x[j+2]==1 and x[j+3]==0:
                    x_mod += [3./np.sqrt(10)-1j*1./np.sqrt(10)]
                elif x[j]==0 and x[j+1]==1 and x[j+2]==1 and x[j+3]==1:
                    x_mod += [3./np.sqrt(10)-1j*3./np.sqrt(10)]
                elif x[j]==1 and x[j+1]==0 and x[j+2]==0 and x[j+3]==0:
                    x_mod += [-1./np.sqrt(10)+1j*1./np.sqrt(10)]
                elif x[j]==1 and x[j+1]==0 and x[j+2]==0 and x[j+3]==1:
                    x_mod += [-1./np.sqrt(10)+1j*3./np.sqrt(10)]
                elif x[j]==1 and x[j+1]==0 and x[j+2]==1 and x[j+3]==0:
                    x_mod += [-3./np.sqrt(10)+1j*1./np.sqrt(10)]
                elif x[j]==1 and x[j+1]==0 and x[j+2]==1 and x[j+3]==1:
                    x_mod += [-3./np.sqrt(10)+1j*3./np.sqrt(10)]
                elif x[j]==1 and x[j+1]==1 and x[j+2]==0 and x[j+3]==0:
                    x_mod += [-1./np.sqrt(10)-1j*1./np.sqrt(10)]
                elif x[j]==1 and x[j+1]==1 and x[j+2]==0 and x[j+3]==1:
                    x_mod += [-1./np.sqrt(10)-1j*3./np.sqrt(10)]
                elif x[j]==1 and x[j+1]==1 and x[j+2]==1 and x[j+3]==0:
                    x_mod += [-3./np.sqrt(10)-1j*1./np.sqrt(10)]
                elif x[j]==1 and x[j+1]==1 and x[j+2]==1 and x[j+3]==1:
                    x_mod += [-3./np.sqrt(10)-1j*3./np.sqrt(10)]
        
        return x_mod

    def scMap(self,x,shift):
        x = list(x)
        Xscm = [0]*(self.fftlen)
        nx = len(x)/12
        if nx % 2 == 0:
            """Even Number"""
            i = len(x)/2
            Xscm[0:i] = x[0:i]
            Xscm[-i::]= x[-i::]
            if shift == True:
                Xscm[-i::] = np.array(Xscm[-i::])*\
                        np.exp(1j*np.pi*np.arange(0,i)/self.fftlen)
                Xscm[0:i] = np.array(Xscm[0:i])*\
                        np.exp(1j*np.pi*np.arange(i,i*2)/self.fftlen)

        else:
            """Odd Number"""
            i = int(np.ceil(len(x)/12/float(2)))*12
            i2 = len(x)-i
            Xscm[-i::] = x[0:i]
            if i2 != 0:
                Xscm[0:i2] = x[-i2::]
            if shift == True:
                Xscm[-i::] = np.array(Xscm[-i::])*np.exp(1j*\
                        np.pi*np.arange(0,i)/self.fftlen)
                if i2 != 0:
                    Xscm[0:i2] = np.array(Xscm[0:i2])*np.exp(1j*\
                           np.pi*np.arange(i,i+i2)/self.fftlen)

        return Xscm

    def scDeMap(self,x,shift):
        x = list(x)
        Xscm = [0]*(self.RB*12)
        if self.RB % 2 == 0:
            """Even Number"""
            i = self.RB/2*12
            Xscm[0:i] = x[0:i]
            Xscm[-i::]= x[-i::]
            if shift == True:
                Xscm[-i::] = np.array(Xscm[-i::])*\
                        np.exp(-1j*np.pi*np.arange(0,i)/self.fftlen)
                Xscm[0:i] = np.array(Xscm[0:i])*\
                        np.exp(-1j*np.pi*np.arange(i,i*2)/self.fftlen)

        else:
            """Odd Number"""
            i = int(np.ceil(self.RB/float(2)))*12
            i2 = self.RB*12-i
            Xscm[0:i] = x[-i::]
            if i2 != 0:
                Xscm[-i2::] = x[0:i2]
            
            if shift == True:
                if i2 != 0:
                    Xscm[0:i2] = np.array(Xscm[0:i2])*np.exp(-1j*\
                            np.pi*np.arange(0,i2)/self.fftlen)
                Xscm[-i::] = np.array(Xscm[-i::])*np.exp(-1j*\
                        np.pi*np.arange(i2,i2+i)/self.fftlen)
        return Xscm

    def addCP(self,x,FSN=False):
        """ FSN is true if it is the first symbol number and else false.
        """
        x = list(x)
        if FSN == True:
            xcp = x[-self.cps0::]+x
        else:
            xcp = x[-self.cps16::]+x
        return xcp

    def remCP(self,x,FSN=False):
        """ FSN is true if it is the first symbol number and else false.
        """
        x = list(x)
        if FSN == True:
            xcp = x[self.cps0::]
        else:
            xcp = x[self.cps16::]
        return xcp

    def ofdm(self,x,symbols,CP,mod,modout=False,shift=True):
        if mod == 'QPSK':
            bits = 2
        else:
            bits = 4
        
        xl = int(len(x)/symbols/bits)
        xtemp = self.modulator(x,mod)

        xcp_all = []
        for i in range(symbols):
            Xmod = xtemp[i*xl:(i+1)*xl]
            
            Xscm = self.scMap(Xmod,shift)
            
            xscm = npf.ifft(Xscm)
            
            if CP == True:
                if i % 7 == 0:
                    xcp = self.addCP(xscm,True)
                else:
                    xcp = self.addCP(xscm)
                xcp_all += xcp
            else:
                xcp_all += list(xscm)
        if modout == True:
            return list(xcp_all),xtemp
        else:
            return list(xcp_all)

    def ofdm_dem(self,x,symbols,CP,shift=True):
        xdem = []
        boff = 0
        for i in range(symbols):
            if CP == True:
                if i % 7 == 0:
                    xl0 = self.fftlen+self.cps0
                    xsym = x[boff:boff+xl0]
                    xcp = self.remCP(xsym,True)
                    boff = boff + xl0
                else:
                    xl16 = self.fftlen+self.cps16
                    xsym =  x[boff:boff+xl16]
                    xcp = self.remCP(xsym)
                    boff = boff + xl16
            else:
                xl = len(x)/symbols
                xcp = x[i*xl:(i+1)*xl]
                
            Xcp = npf.fft(xcp)
            
            Xsc = self.scDeMap(Xcp,shift)
            
            xdem += list(Xsc)
        return xdem

    def sc_fdm(self,x,symbols,CP,mod,modout=False,shift=True):
        if mod == 'QPSK':
            bits = 2
        else:
            bits = 4

        xl = int(len(x)/symbols/bits)
        xtemp = self.modulator(x,mod)

        xcp_all = []
        for i in range(symbols):
            Xmod = xtemp[i*xl:(i+1)*xl]

            if mod == '16QAM':
                bits = 4.
            else:
                bits = 2.

            Xdft = 1./(np.sqrt(len(x)/symbols/bits))*npf.fft(Xmod)
            
            Xscm = self.scMap(Xdft,shift)
            
            xscm = npf.ifft(Xscm)
            
            if CP == True:
                if i % 7 == 0:
                    xcp = self.addCP(xscm,True)
                else:
                    xcp = self.addCP(xscm)
                xcp_all += xcp
            else:
                xcp_all += list(xscm)
        if modout == True:
            return list(xcp_all),xtemp
        else:
            return list(xcp_all)
    
    def sc_fdm_dem(self,x,symbols,CP,shift=True):
        xdem = []
        boff = 0
        for i in range(symbols):
            if CP == True:
                if i % 7 == 0:
                    xl0 = self.fftlen+self.cps0
                    xsym = x[boff:boff+xl0]
                    xcp = self.remCP(xsym,True)
                    boff = boff + xl0
                else:
                    xl16 = self.fftlen+self.cps16
                    xsym =  x[boff:boff+xl16]
                    xcp = self.remCP(xsym)
                    boff = boff + xl16
            else:
                xl = len(x)/symbols
                xcp = x[i*xl:(i+1)*xl]
            
            Xcp = npf.fft(xcp)
            
            Xsc = self.scDeMap(Xcp,shift)
            
            xdft = np.sqrt(len(Xsc))*npf.ifft(Xsc)

            xdem += list(xdft)
        return xdem

        
class RFFE:
    def __init__(self,fc,fftlen):
        self.fc = fc
        self.fftlen = fftlen
        self.scs = 15E3
        self.Fs = self.fftlen/(1./self.scs)
        
        self.ps = misc.PulseShaping(self.Fs)

    def plotTime(self):
        return np.arange(0,1./(self.Fs)*self.xl,1./(self.Fs*self.osamp))

    def plotFreq(self):
        return np.fft.fftfreq(self.xlos, 1./(self.Fs*self.osamp))

    def dac(self,x,OS=4):
        """
        Resamples x to fit the carrier frequency
        """

        self.Fsnr = int(round(self.fc*OS/self.Fs)*len(x))
        self.osamp = self.Fsnr/len(x)
        xos = list(sps.resample(x,self.Fsnr))
        
        self.xl = len(x)
        self.xlos = len(xos)
        return xos

    def adc(self,x,xlen=False):
        if xlen == False:
            xlen = self.xl
        return sps.resample(x,xlen)

    def pureSinCos(self):
        pure = [1E-20 for i in range(self.xlos)]
        i_temp = int(round(self.fc/(self.Fs*self.osamp/self.xlos)))

        pure[i_temp] = self.xlos
        cos = np.real(np.fft.ifft(pure))
        sin = np.imag(np.fft.ifft(pure))
        return sin,cos

    def upconversion(self,x):
        Irf = np.real(x)*self.pureSinCos()[1]
        Qrf = np.imag(x)*-1*self.pureSinCos()[0]
        return Irf + Qrf

    def downconversion(self,x,fco):
        sin,cos = self.pureSinCos()
        
        Id = x*2*cos
        Idf = self.ps.rect(Id,fco)

        Qd = x*-2*sin
        Qdf = self.ps.rect(Qd,fco)

        return Idf+1j*Qdf

    def tx(self,x):
        xt = self.dac(x)
        return self.upconversion(xt)

    def rx(self,x,fco):
        yr = self.downconversion(x,fco)
        return self.adc(yr)


